import threading
import time
from concurrent.futures import ThreadPoolExecutor,wait
from sys import argv
import requests
from urllib.parse import urljoin
import os

ts_list = []
start_time = None
left_cnt = 0
all_sz = 0

def to_mp4():
    print("[+] 合成中...")
    file_name ="./videos/%s.mp4"%start_time
    with open(file_name,'wb') as mp4:
        for i in range(0,len(ts_list)):
            with open("./videos/%s/%d.ts"%(start_time,i),'rb') as ts:
                mp4.write(ts.read())
            os.remove("./videos/%s/%d.ts"%(start_time,i))
    print("[+] 合成完毕，总用时：%d秒，路径：%s"%((time.time()-start_time),file_name))


def download_one(index):
    global left_cnt,all_sz
    url = ts_list[index]

    print("[+] 开始下载任务%d"%index)
    resp = requests.get(url=url)
    print("[+] 任务%d，完成，写入中..."%index)

    if resp.status_code!=200:
        print("[-] 任务%d下载失败"%index)

    ts_size = len(resp.content)/1000/1024
    all_sz=ts_size+all_sz
    v=all_sz/((time.time()-start_time))


    with open("./videos/%s/%d.ts"%(start_time,index),'wb') as f:
        for chunk in resp.iter_content(1024):
            f.write(chunk)
    left_cnt=left_cnt-1
    print("[+] 任务%d，写入完成，剩余任务数量%d，速度%.2fM/s"%(index,left_cnt,v))
    if left_cnt==0:
        to_mp4()



def main():
    max_thread_cnt = 100

    v=False
    try:
        max_thread_cnt = int(argv[2])
        v=argv[3]
        v=True
    except:
        pass

    print("[+] 最大线程数 %d"%max_thread_cnt)
    global ts_list,start_time,left_cnt
    print("[+] 正在解析m3u8...")
    ts_resp = requests.get(argv[1])
    if ts_resp.status_code != 200:
        print("[-] 返回了错误的代码：%d"%ts_resp.status_code)
    ts_list = [urljoin(argv[1],ts) for ts in ts_resp.text.split("\n") if ts!="" and ts[0]!="#"]
    if v:
        print(ts_list)
    left_cnt=len(ts_list)
    print("[+] 共%d个任务"%left_cnt)

    start_time = time.time()
    os.mkdir("./videos/%s"%start_time)
    with ThreadPoolExecutor(max_thread_cnt) as excecutor:
        tasks = [excecutor.submit(download_one,i) for i in range(len(ts_list))]
        ret_tasks = wait(tasks,timeout=200)
        print(ret_tasks)


if __name__ == '__main__':
    main()